{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 从离散动作空间到连续动作空间的VPG算法\n",
    "\n",
    "连续动作空间下的VPG算法拟合的是一个分布，可以简单地使用高斯分布（需要说明的是，**一个高斯分布经过足够复杂的映射可以拟合任意复杂的分布**）。于是有如下算法：\n",
    "$$\\pi(a_{t}|s_{t};\\theta^{\\mu,\\sigma})=N(\\mu(s_{t}),\\sigma(s_{t})) \\\\ \n",
    "a_{t}\\sim \\pi(A|s_{t};\\theta^{\\mu,\\sigma}) \\\\\n",
    "Loss(\\theta^{\\mu,\\sigma})=Adv(s_{t},a_{t})*log\\pi(a_{t}|s_{t};\\theta^{\\mu,\\sigma})$$\n",
    "VPG算法是一个同策略算法，也就是不使用buffer replay，原因很好理解，**旧有的经验不适用于更新后的策略**，所以只能采样一条轨迹更新一次参数的进行梯度更新。\n",
    "\n",
    "#### 从sample()到rsample()\n",
    "\n",
    "sample()函数从一个分布中抽样一个动作，rsample()函数则使用重参数化技巧：\n",
    "$$a\\sim N(\\mu(s_{t}),\\sigma(s_{t})) \\\\ a=\\mu(s_{t})+\\sigma(s_{t})\\odot \\xi,\\xi\\sim N(0,1)$$\n",
    "使用重参数化技巧的原因是**计算梯度时动作a需要对参数$\\theta^{\\mu,\\sigma}$求导**。\n",
    "\n",
    "#### 同策略算法对比异策略算法\n",
    "\n",
    "可以看到，VPG算法不需要a对参数求导而SAC需要，其原因在于：\n",
    "1. VPG算法是同策略算法，采样出来的动作不具有滞后性，a作为一个标量直接送入损失函数中即可；\n",
    "2. SAC算法是异策略算法，采样出来的动作作为经验，作用是求评论家(评估)的损失值，让评论家(评估)越来越接近ground truth。损失函数中的动作是通过**当前**策略采样得到的，因此求导时动作必须对参数求导。\n",
    "\n",
    "对于连续动作空间下的异策略算法，经验(s,a,r,s')中的动作a较策略而言已经过时，a的作用是提供参数，**让评论家(评估)在状态s、做出动作a的价值预测得更精准**，和更新策略无关。\n",
    "\n",
    "#### SAC算法的核心思想\n",
    "\n",
    "连续动作空间下的策略最终得到的是**一个分布**，这个分布给出了动作范围中**哪些动作有更大的可能性被选中**。动作概率分布可以是单峰的（确定性策略，例如DDPG、TD3），也可以是多峰的（SAC），也可以是其他形状比如均匀的（随机策略）。确定性策略看起来采样出一个确定的动作，实际上还是一个分布，只不过最优动作的概率无限接近于1。\n",
    "\n",
    "SAC算法的核心思想是**在连续动作空间下拟合一个多峰的策略网络**，SAC相较于确定性策略（DDPG、TD3）的优点是可以处理最优解并非只有一种情况的场景，也就是泛用性更强，能够探索更多情况。为此，SAC算法重新定义了状态价值函数：\n",
    "$$V(s)=E_{a_{t}\\sim\\pi}(Q(s_{t},a_{t})-\\alpha log\\pi(a_{t}|s_{t}))$$\n",
    "\n",
    "SAC原始论文通过一系列数学证明，证明了上式能够收敛到最优以及通过策略评估（评论家更新）和策略迭代（演员更新）这两个过程交替进行能够得到最优策略。数学证明过程看似眼花缭乱，实际最后还是**通过梯度上升最大化状态价值函数V(s)**。\n",
    "\n",
    "#### SAC算法的训练参数\n",
    "\n",
    "SAC算法拟合了一个高斯分布作为策略，因为是异策略算法，损失函数中的动作需要通过当前策略采样，因此需要对动作进行重参数化采样：\n",
    "$$\\pi(a_{t}|s_{t};\\theta^{\\mu},\\theta^{\\sigma})=N(\\mu(s_{t}),\\sigma(s_{t})) \\\\ \n",
    "a=tanh(\\mu(s_{t})+\\sigma(s_{t})\\odot\\xi),\\xi\\sim N(0,1)$$\n",
    "\n",
    "#### SAC算法的模型及参数\n",
    "\n",
    "1. 一组经验(s,a,r,s')\n",
    "2. 演员$\\pi(\\cdot|s;\\theta^{\\mu},\\theta^{\\sigma})$\n",
    "3. 评论家（评估）$Q_{1}$、$Q_{2}$；评论家（目标）$Q_{1}^{'}$、$Q_{2}^{'}$\n",
    "\n",
    "两个评论家(目标)的原因是为了取较小值抑制Q值高估问题。\n",
    "\n",
    "core.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.signal\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.distributions.normal import Normal\n",
    "\n",
    "\n",
    "def combined_shape(length, shape=None):\n",
    "    if shape is None:\n",
    "        return (length,)\n",
    "    return (length, shape) if np.isscalar(shape) else (length, *shape)\n",
    "\n",
    "def mlp(sizes, activation, output_activation=nn.Identity):\n",
    "    layers = []\n",
    "    for j in range(len(sizes)-1):\n",
    "        act = activation if j < len(sizes)-2 else output_activation\n",
    "        layers += [nn.Linear(sizes[j], sizes[j+1]), act()]\n",
    "    return nn.Sequential(*layers)\n",
    "\n",
    "def count_vars(module):\n",
    "    return sum([np.prod(p.shape) for p in module.parameters()])\n",
    "\n",
    "\n",
    "LOG_STD_MAX = 2\n",
    "LOG_STD_MIN = -20\n",
    "\n",
    "class SquashedGaussianMLPActor(nn.Module):\n",
    "\n",
    "    def __init__(self, obs_dim, act_dim, hidden_sizes, activation, act_limit):\n",
    "        super().__init__()\n",
    "        self.net = mlp([obs_dim] + list(hidden_sizes), activation, activation)\n",
    "        self.mu_layer = nn.Linear(hidden_sizes[-1], act_dim)\n",
    "        self.log_std_layer = nn.Linear(hidden_sizes[-1], act_dim)\n",
    "        self.act_limit = act_limit\n",
    "\n",
    "    def forward(self, obs, deterministic=False, with_logprob=True):\n",
    "        net_out = self.net(obs)\n",
    "        mu = self.mu_layer(net_out)\n",
    "        log_std = self.log_std_layer(net_out)\n",
    "        log_std = torch.clamp(log_std, LOG_STD_MIN, LOG_STD_MAX)\n",
    "        std = torch.exp(log_std)\n",
    "\n",
    "        # Pre-squash distribution and sample\n",
    "        pi_distribution = Normal(mu, std)\n",
    "        if deterministic:\n",
    "            # Only used for evaluating policy at test time.\n",
    "            pi_action = mu\n",
    "        else:\n",
    "            pi_action = pi_distribution.rsample()\n",
    "\n",
    "        if with_logprob:\n",
    "            # Compute logprob from Gaussian, and then apply correction for Tanh squashing.\n",
    "            # NOTE: The correction formula is a little bit magic. To get an understanding \n",
    "            # of where it comes from, check out the original SAC paper (arXiv 1801.01290) \n",
    "            # and look in appendix C. This is a more numerically-stable equivalent to Eq 21.\n",
    "            # Try deriving it yourself as a (very difficult) exercise. :)\n",
    "            logp_pi = pi_distribution.log_prob(pi_action).sum(axis=-1)\n",
    "            logp_pi -= (2*(np.log(2) - pi_action - F.softplus(-2*pi_action))).sum(axis=1)\n",
    "        else:\n",
    "            logp_pi = None\n",
    "\n",
    "        pi_action = torch.tanh(pi_action)\n",
    "        pi_action = self.act_limit * pi_action\n",
    "\n",
    "        return pi_action, logp_pi\n",
    "\n",
    "\n",
    "class MLPQFunction(nn.Module):\n",
    "\n",
    "    def __init__(self, obs_dim, act_dim, hidden_sizes, activation):\n",
    "        super().__init__()\n",
    "        self.q = mlp([obs_dim + act_dim] + list(hidden_sizes) + [1], activation)\n",
    "\n",
    "    def forward(self, obs, act):\n",
    "        q = self.q(torch.cat([obs, act], dim=-1))\n",
    "        return torch.squeeze(q, -1) # Critical to ensure q has right shape.\n",
    "\n",
    "class MLPActorCritic(nn.Module):\n",
    "\n",
    "    def __init__(self, observation_space, action_space, hidden_sizes=(256,256),\n",
    "                 activation=nn.ReLU):\n",
    "        super().__init__()\n",
    "\n",
    "        obs_dim = observation_space.shape[0]\n",
    "        act_dim = action_space.shape[0]\n",
    "        act_limit = action_space.high[0]\n",
    "\n",
    "        # build policy and value functions\n",
    "        self.pi = SquashedGaussianMLPActor(obs_dim, act_dim, hidden_sizes, activation, act_limit)\n",
    "        self.q1 = MLPQFunction(obs_dim, act_dim, hidden_sizes, activation)\n",
    "        self.q2 = MLPQFunction(obs_dim, act_dim, hidden_sizes, activation)\n",
    "\n",
    "    def act(self, obs, deterministic=False):\n",
    "        with torch.no_grad():\n",
    "            a, _ = self.pi(obs, deterministic, False)\n",
    "            return a.numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "sac.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy\n",
    "import itertools\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.optim import Adam\n",
    "import gym\n",
    "import time\n",
    "import spinup.algos.pytorch.sac.core as core\n",
    "from spinup.utils.logx import EpochLogger\n",
    "\n",
    "class ReplayBuffer:\n",
    "    \"\"\"\n",
    "    A simple FIFO experience replay buffer for SAC agents.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, obs_dim, act_dim, size):\n",
    "        self.obs_buf = np.zeros(core.combined_shape(size, obs_dim), dtype=np.float32)\n",
    "        self.obs2_buf = np.zeros(core.combined_shape(size, obs_dim), dtype=np.float32)\n",
    "        self.act_buf = np.zeros(core.combined_shape(size, act_dim), dtype=np.float32)\n",
    "        self.rew_buf = np.zeros(size, dtype=np.float32)\n",
    "        self.done_buf = np.zeros(size, dtype=np.float32)\n",
    "        self.ptr, self.size, self.max_size = 0, 0, size\n",
    "\n",
    "    def store(self, obs, act, rew, next_obs, done):\n",
    "        self.obs_buf[self.ptr] = obs\n",
    "        self.obs2_buf[self.ptr] = next_obs\n",
    "        self.act_buf[self.ptr] = act\n",
    "        self.rew_buf[self.ptr] = rew\n",
    "        self.done_buf[self.ptr] = done\n",
    "        self.ptr = (self.ptr+1) % self.max_size\n",
    "        self.size = min(self.size+1, self.max_size)\n",
    "\n",
    "    def sample_batch(self, batch_size=32):\n",
    "        idxs = np.random.randint(0, self.size, size=batch_size)\n",
    "        batch = dict(obs=self.obs_buf[idxs],\n",
    "                     obs2=self.obs2_buf[idxs],\n",
    "                     act=self.act_buf[idxs],\n",
    "                     rew=self.rew_buf[idxs],\n",
    "                     done=self.done_buf[idxs])\n",
    "        return {k: torch.as_tensor(v, dtype=torch.float32) for k,v in batch.items()}\n",
    "\n",
    "def sac(env_fn, actor_critic=core.MLPActorCritic, ac_kwargs=dict(), seed=0, \n",
    "        steps_per_epoch=4000, epochs=100, replay_size=int(1e6), gamma=0.99, \n",
    "        polyak=0.995, lr=1e-3, alpha=0.2, batch_size=100, start_steps=10000, \n",
    "        update_after=1000, update_every=50, num_test_episodes=10, max_ep_len=1000, \n",
    "        logger_kwargs=dict(), save_freq=1):\n",
    "    \"\"\"\n",
    "    Soft Actor-Critic (SAC)\n",
    "\n",
    "\n",
    "    Args:\n",
    "        env_fn : A function which creates a copy of the environment.\n",
    "            The environment must satisfy the OpenAI Gym API.\n",
    "\n",
    "        actor_critic: The constructor method for a PyTorch Module with an ``act`` \n",
    "            method, a ``pi`` module, a ``q1`` module, and a ``q2`` module.\n",
    "            The ``act`` method and ``pi`` module should accept batches of \n",
    "            observations as inputs, and ``q1`` and ``q2`` should accept a batch \n",
    "            of observations and a batch of actions as inputs. When called, \n",
    "            ``act``, ``q1``, and ``q2`` should return:\n",
    "\n",
    "            ===========  ================  ======================================\n",
    "            Call         Output Shape      Description\n",
    "            ===========  ================  ======================================\n",
    "            ``act``      (batch, act_dim)  | Numpy array of actions for each \n",
    "                                           | observation.\n",
    "            ``q1``       (batch,)          | Tensor containing one current estimate\n",
    "                                           | of Q* for the provided observations\n",
    "                                           | and actions. (Critical: make sure to\n",
    "                                           | flatten this!)\n",
    "            ``q2``       (batch,)          | Tensor containing the other current \n",
    "                                           | estimate of Q* for the provided observations\n",
    "                                           | and actions. (Critical: make sure to\n",
    "                                           | flatten this!)\n",
    "            ===========  ================  ======================================\n",
    "\n",
    "            Calling ``pi`` should return:\n",
    "\n",
    "            ===========  ================  ======================================\n",
    "            Symbol       Shape             Description\n",
    "            ===========  ================  ======================================\n",
    "            ``a``        (batch, act_dim)  | Tensor containing actions from policy\n",
    "                                           | given observations.\n",
    "            ``logp_pi``  (batch,)          | Tensor containing log probabilities of\n",
    "                                           | actions in ``a``. Importantly: gradients\n",
    "                                           | should be able to flow back into ``a``.\n",
    "            ===========  ================  ======================================\n",
    "\n",
    "        ac_kwargs (dict): Any kwargs appropriate for the ActorCritic object \n",
    "            you provided to SAC.\n",
    "\n",
    "        seed (int): Seed for random number generators.\n",
    "\n",
    "        steps_per_epoch (int): Number of steps of interaction (state-action pairs) \n",
    "            for the agent and the environment in each epoch.\n",
    "\n",
    "        epochs (int): Number of epochs to run and train agent.\n",
    "\n",
    "        replay_size (int): Maximum length of replay buffer.\n",
    "\n",
    "        gamma (float): Discount factor. (Always between 0 and 1.)\n",
    "\n",
    "        polyak (float): Interpolation factor in polyak averaging for target \n",
    "            networks. Target networks are updated towards main networks \n",
    "            according to:\n",
    "\n",
    "            .. math:: \\\\theta_{\\\\text{targ}} \\\\leftarrow \n",
    "                \\\\rho \\\\theta_{\\\\text{targ}} + (1-\\\\rho) \\\\theta\n",
    "\n",
    "            where :math:`\\\\rho` is polyak. (Always between 0 and 1, usually \n",
    "            close to 1.)\n",
    "\n",
    "        lr (float): Learning rate (used for both policy and value learning).\n",
    "\n",
    "        alpha (float): Entropy regularization coefficient. (Equivalent to \n",
    "            inverse of reward scale in the original SAC paper.)\n",
    "\n",
    "        batch_size (int): Minibatch size for SGD.\n",
    "\n",
    "        start_steps (int): Number of steps for uniform-random action selection,\n",
    "            before running real policy. Helps exploration.\n",
    "\n",
    "        update_after (int): Number of env interactions to collect before\n",
    "            starting to do gradient descent updates. Ensures replay buffer\n",
    "            is full enough for useful updates.\n",
    "\n",
    "        update_every (int): Number of env interactions that should elapse\n",
    "            between gradient descent updates. Note: Regardless of how long \n",
    "            you wait between updates, the ratio of env steps to gradient steps \n",
    "            is locked to 1.\n",
    "\n",
    "        num_test_episodes (int): Number of episodes to test the deterministic\n",
    "            policy at the end of each epoch.\n",
    "\n",
    "        max_ep_len (int): Maximum length of trajectory / episode / rollout.\n",
    "\n",
    "        logger_kwargs (dict): Keyword args for EpochLogger.\n",
    "\n",
    "        save_freq (int): How often (in terms of gap between epochs) to save\n",
    "            the current policy and value function.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    logger = EpochLogger(**logger_kwargs)\n",
    "    logger.save_config(locals())\n",
    "\n",
    "    torch.manual_seed(seed)\n",
    "    np.random.seed(seed)\n",
    "\n",
    "    env, test_env = env_fn(), env_fn()\n",
    "    obs_dim = env.observation_space.shape\n",
    "    act_dim = env.action_space.shape[0]\n",
    "\n",
    "    # Action limit for clamping: critically, assumes all dimensions share the same bound!\n",
    "    act_limit = env.action_space.high[0]\n",
    "\n",
    "    # Create actor-critic module and target networks\n",
    "    ac = actor_critic(env.observation_space, env.action_space, **ac_kwargs)\n",
    "    ac_targ = deepcopy(ac)\n",
    "\n",
    "    # Freeze target networks with respect to optimizers (only update via polyak averaging)\n",
    "    for p in ac_targ.parameters():\n",
    "        p.requires_grad = False\n",
    "        \n",
    "    # List of parameters for both Q-networks (save this for convenience)\n",
    "    q_params = itertools.chain(ac.q1.parameters(), ac.q2.parameters())\n",
    "\n",
    "    # Experience buffer\n",
    "    replay_buffer = ReplayBuffer(obs_dim=obs_dim, act_dim=act_dim, size=replay_size)\n",
    "\n",
    "    # Count variables (protip: try to get a feel for how different size networks behave!)\n",
    "    var_counts = tuple(core.count_vars(module) for module in [ac.pi, ac.q1, ac.q2])\n",
    "    logger.log('\\nNumber of parameters: \\t pi: %d, \\t q1: %d, \\t q2: %d\\n'%var_counts)\n",
    "\n",
    "    # Set up function for computing SAC Q-losses\n",
    "    def compute_loss_q(data):\n",
    "        o, a, r, o2, d = data['obs'], data['act'], data['rew'], data['obs2'], data['done']\n",
    "\n",
    "        q1 = ac.q1(o,a)\n",
    "        q2 = ac.q2(o,a)\n",
    "\n",
    "        # Bellman backup for Q functions\n",
    "        with torch.no_grad():\n",
    "            # Target actions come from *current* policy\n",
    "            a2, logp_a2 = ac.pi(o2)\n",
    "\n",
    "            # Target Q-values\n",
    "            q1_pi_targ = ac_targ.q1(o2, a2)\n",
    "            q2_pi_targ = ac_targ.q2(o2, a2)\n",
    "            q_pi_targ = torch.min(q1_pi_targ, q2_pi_targ)\n",
    "            backup = r + gamma * (1 - d) * (q_pi_targ - alpha * logp_a2)\n",
    "\n",
    "        # MSE loss against Bellman backup\n",
    "        loss_q1 = ((q1 - backup)**2).mean()\n",
    "        loss_q2 = ((q2 - backup)**2).mean()\n",
    "        loss_q = loss_q1 + loss_q2\n",
    "\n",
    "        # Useful info for logging\n",
    "        q_info = dict(Q1Vals=q1.detach().numpy(),\n",
    "                      Q2Vals=q2.detach().numpy())\n",
    "\n",
    "        return loss_q, q_info\n",
    "\n",
    "    # Set up function for computing SAC pi loss\n",
    "    def compute_loss_pi(data):\n",
    "        o = data['obs']\n",
    "        pi, logp_pi = ac.pi(o)\n",
    "        q1_pi = ac.q1(o, pi)\n",
    "        q2_pi = ac.q2(o, pi)\n",
    "        q_pi = torch.min(q1_pi, q2_pi)\n",
    "\n",
    "        # Entropy-regularized policy loss\n",
    "        loss_pi = (alpha * logp_pi - q_pi).mean()\n",
    "\n",
    "        # Useful info for logging\n",
    "        pi_info = dict(LogPi=logp_pi.detach().numpy())\n",
    "\n",
    "        return loss_pi, pi_info\n",
    "\n",
    "    # Set up optimizers for policy and q-function\n",
    "    pi_optimizer = Adam(ac.pi.parameters(), lr=lr)\n",
    "    q_optimizer = Adam(q_params, lr=lr)\n",
    "\n",
    "    # Set up model saving\n",
    "    logger.setup_pytorch_saver(ac)\n",
    "\n",
    "    def update(data):\n",
    "        # First run one gradient descent step for Q1 and Q2\n",
    "        q_optimizer.zero_grad()\n",
    "        loss_q, q_info = compute_loss_q(data)\n",
    "        loss_q.backward()\n",
    "        q_optimizer.step()\n",
    "\n",
    "        # Record things\n",
    "        logger.store(LossQ=loss_q.item(), **q_info)\n",
    "\n",
    "        # Freeze Q-networks so you don't waste computational effort \n",
    "        # computing gradients for them during the policy learning step.\n",
    "        for p in q_params:\n",
    "            p.requires_grad = False\n",
    "\n",
    "        # Next run one gradient descent step for pi.\n",
    "        pi_optimizer.zero_grad()\n",
    "        loss_pi, pi_info = compute_loss_pi(data)\n",
    "        loss_pi.backward()\n",
    "        pi_optimizer.step()\n",
    "\n",
    "        # Unfreeze Q-networks so you can optimize it at next DDPG step.\n",
    "        for p in q_params:\n",
    "            p.requires_grad = True\n",
    "\n",
    "        # Record things\n",
    "        logger.store(LossPi=loss_pi.item(), **pi_info)\n",
    "\n",
    "        # Finally, update target networks by polyak averaging.\n",
    "        with torch.no_grad():\n",
    "            for p, p_targ in zip(ac.parameters(), ac_targ.parameters()):\n",
    "                # NB: We use an in-place operations \"mul_\", \"add_\" to update target\n",
    "                # params, as opposed to \"mul\" and \"add\", which would make new tensors.\n",
    "                p_targ.data.mul_(polyak)\n",
    "                p_targ.data.add_((1 - polyak) * p.data)\n",
    "\n",
    "    def get_action(o, deterministic=False):\n",
    "        return ac.act(torch.as_tensor(o, dtype=torch.float32), \n",
    "                      deterministic)\n",
    "\n",
    "    def test_agent():\n",
    "        for j in range(num_test_episodes):\n",
    "            o, d, ep_ret, ep_len = test_env.reset(), False, 0, 0\n",
    "            while not(d or (ep_len == max_ep_len)):\n",
    "                # Take deterministic actions at test time \n",
    "                o, r, d, _ = test_env.step(get_action(o, True))\n",
    "                ep_ret += r\n",
    "                ep_len += 1\n",
    "            logger.store(TestEpRet=ep_ret, TestEpLen=ep_len)\n",
    "\n",
    "    # Prepare for interaction with environment\n",
    "    total_steps = steps_per_epoch * epochs\n",
    "    start_time = time.time()\n",
    "    o, ep_ret, ep_len = env.reset(), 0, 0\n",
    "\n",
    "    # Main loop: collect experience in env and update/log each epoch\n",
    "    for t in range(total_steps):\n",
    "        \n",
    "        # Until start_steps have elapsed, randomly sample actions\n",
    "        # from a uniform distribution for better exploration. Afterwards, \n",
    "        # use the learned policy. \n",
    "        if t > start_steps:\n",
    "            a = get_action(o)\n",
    "        else:\n",
    "            a = env.action_space.sample()\n",
    "\n",
    "        # Step the env\n",
    "        o2, r, d, _ = env.step(a)\n",
    "        ep_ret += r\n",
    "        ep_len += 1\n",
    "\n",
    "        # Ignore the \"done\" signal if it comes from hitting the time\n",
    "        # horizon (that is, when it's an artificial terminal signal\n",
    "        # that isn't based on the agent's state)\n",
    "        d = False if ep_len==max_ep_len else d\n",
    "\n",
    "        # Store experience to replay buffer\n",
    "        replay_buffer.store(o, a, r, o2, d)\n",
    "\n",
    "        # Super critical, easy to overlook step: make sure to update \n",
    "        # most recent observation!\n",
    "        o = o2\n",
    "\n",
    "        # End of trajectory handling\n",
    "        if d or (ep_len == max_ep_len):\n",
    "            logger.store(EpRet=ep_ret, EpLen=ep_len)\n",
    "            o, ep_ret, ep_len = env.reset(), 0, 0\n",
    "\n",
    "        # Update handling\n",
    "        if t >= update_after and t % update_every == 0:\n",
    "            for j in range(update_every):\n",
    "                batch = replay_buffer.sample_batch(batch_size)\n",
    "                update(data=batch)\n",
    "\n",
    "        # End of epoch handling\n",
    "        if (t+1) % steps_per_epoch == 0:\n",
    "            epoch = (t+1) // steps_per_epoch\n",
    "\n",
    "            # Save model\n",
    "            if (epoch % save_freq == 0) or (epoch == epochs):\n",
    "                logger.save_state({'env': env}, None)\n",
    "\n",
    "            # Test the performance of the deterministic version of the agent.\n",
    "            test_agent()\n",
    "\n",
    "            # Log info about epoch\n",
    "            logger.log_tabular('Epoch', epoch)\n",
    "            logger.log_tabular('EpRet', with_min_and_max=True)\n",
    "            logger.log_tabular('TestEpRet', with_min_and_max=True)\n",
    "            logger.log_tabular('EpLen', average_only=True)\n",
    "            logger.log_tabular('TestEpLen', average_only=True)\n",
    "            logger.log_tabular('TotalEnvInteracts', t)\n",
    "            logger.log_tabular('Q1Vals', with_min_and_max=True)\n",
    "            logger.log_tabular('Q2Vals', with_min_and_max=True)\n",
    "            logger.log_tabular('LogPi', with_min_and_max=True)\n",
    "            logger.log_tabular('LossPi', average_only=True)\n",
    "            logger.log_tabular('LossQ', average_only=True)\n",
    "            logger.log_tabular('Time', time.time()-start_time)\n",
    "            logger.dump_tabular()\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    import argparse\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument('--env', type=str, default='HalfCheetah-v2')\n",
    "    parser.add_argument('--hid', type=int, default=256)\n",
    "    parser.add_argument('--l', type=int, default=2)\n",
    "    parser.add_argument('--gamma', type=float, default=0.99)\n",
    "    parser.add_argument('--seed', '-s', type=int, default=0)\n",
    "    parser.add_argument('--epochs', type=int, default=50)\n",
    "    parser.add_argument('--exp_name', type=str, default='sac')\n",
    "    args = parser.parse_args()\n",
    "\n",
    "    from spinup.utils.run_utils import setup_logger_kwargs\n",
    "    logger_kwargs = setup_logger_kwargs(args.exp_name, args.seed)\n",
    "\n",
    "    torch.set_num_threads(torch.get_num_threads())\n",
    "\n",
    "    sac(lambda : gym.make(args.env), actor_critic=core.MLPActorCritic,\n",
    "        ac_kwargs=dict(hidden_sizes=[args.hid]*args.l), \n",
    "        gamma=args.gamma, seed=args.seed, epochs=args.epochs,\n",
    "        logger_kwargs=logger_kwargs)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
